package rh.cola.api;

import com.unfbx.chatgpt.OpenAiClient;
import com.unfbx.chatgpt.OpenAiStreamClient;
import com.unfbx.chatgpt.entity.chat.ChatCompletion;
import com.unfbx.chatgpt.entity.chat.Message;
import com.unfbx.chatgpt.entity.completions.Completion;
import com.unfbx.chatgpt.entity.completions.CompletionResponse;
import com.unfbx.chatgpt.sse.ConsoleEventSourceListener;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.Test;
import org.springframework.boot.test.context.SpringBootTest;
import rh.cola.utils.ChatGPTUtils;

import java.util.Arrays;
import java.util.concurrent.CountDownLatch;

/**
 * @Name: HelloChatGPT
 * @Author: Cola
 * @Time: 2023/3/22 16:55
 * @Description: HelloChatGPT
 */
@SpringBootTest
@Slf4j
public class HelloChatGPT {

    /**
     * 常规客户端
     */
    @Test
    void helloGPTForCommonClient() {
        /*获取客户端*/
        OpenAiClient commonClientByApiKey = ChatGPTUtils.getCommonClientByApiKey();
        /*对话*/
        CompletionResponse completions = commonClientByApiKey
                .completions("你叫什么名字");
        Arrays.stream(completions.getChoices()).forEach(System.out::println);
    }

    /**
     * 流式输出
     */
    @Test
    void helloGPTForStreamClientV1() {
        ConsoleEventSourceListener eventSourceListener = new ConsoleEventSourceListener();
        Message message = Message.builder().role(Message.Role.USER).content("你好啊我的伙伴").build();
        ChatCompletion chatCompletion = ChatCompletion.builder().messages(Arrays.asList(message)).build();
        /*获取客户端*/
        OpenAiStreamClient streamClientByApiKey = ChatGPTUtils.getStreamClientByApiKey();
        streamClientByApiKey.streamChatCompletion(chatCompletion, eventSourceListener);
        CountDownLatch countDownLatch = new CountDownLatch(1);
        try {
            countDownLatch.await();
        } catch (InterruptedException e) {
            if (log.isDebugEnabled()) {
                log.debug(e.toString());
            }
        }

    }

    /**
     * 流式输出
     */
    @Test
    void helloGPTForStreamClientV2() {
        ConsoleEventSourceListener eventSourceListener = new ConsoleEventSourceListener();
        Completion completion = Completion.builder()
                .prompt("一句话描述一下现在的心情")
                .stream(true)
                .build();
        /*获取客户端*/
        OpenAiStreamClient streamClientByApiKey = ChatGPTUtils.getStreamClientByApiKey();
        streamClientByApiKey.streamCompletions(completion, eventSourceListener);
        CountDownLatch countDownLatch = new CountDownLatch(1);

        try {
            countDownLatch.await();
        } catch (InterruptedException e) {
            if (log.isDebugEnabled()) {
                log.debug(e.toString());
            }
        }
    }
}
